#!/usr/bin/python
# -*- coding: utf-8 -*-
import pymysql
import time
from urlparse import urlparse
from threading import Lock
from burp import IBurpExtender, ITab
from burp import IHttpListener
from java.io import PrintWriter
from java.util import ArrayList
from javax.swing import BorderFactory
from javax.swing import BoxLayout
from javax.swing import JTextField
from javax.swing import JButton
from java.awt.event import ActionListener
from javax.swing import JPanel
from javax.swing import JTabbedPane


class connectMysql(ActionListener):
    def __init__(self, extender):
        self._extender = extender

    # 连接Mysql数据库
    def actionPerformed(self, e):
        button = self._extender.connect_button
        if button.getText() == "Connect":
            self._extender._mysql = pymysql.connect(host='127.0.0.1',
                                                    port=3306,
                                                    user='dbuser',
                                                    password='Taochao123',
                                                    database='webmonitor',
                                                    charset='utf8')
            button.setText("Disconnect ")
        else:
            button.setText("Connect")
            self._extender._mysql.close()
        return


def insert_url_rec(_mysql, _log):
    con = _mysql
    cur = con.cursor()
    try:
        # 防止sql注入
        # sql_str = "INSERT INTO url(project_id,host, method,url,params) VALUES ('%s','%s','%s','%s','%s')" \
        # % (1,_log._host,_log._method,_log._url,_log.fuzzpayload)

        # sql_str = "INSERT INTO url(project_id,host, method,url,params) VALUES(2,'_log._host','http//:www.bass.com','PUT','file')"
        sql_str = "INSERT INTO url(project_id,host,method,url,params,status,mime) VALUES ('%s','%s','%s','%s','%s','%s','%s')" \
                  % (_log._project_id,_log._host,_log._method,_log._url, _log._query,_log._status,_log._mime)
        cur.execute(sql_str)
        con.commit()
    except Exception as e:
        con.rollback()
        print("connect mysql error", e)
        return


class BurpExtender(IBurpExtender, ITab, IHttpListener):
    # Burp extensions 列表中的扩展名
    _extensionName = "BurpMysql"
    _labelName = "Burp Mysql"
    _mainName = "main"

    def registerExtenderCallbacks(self, callbacks):
        # 表示私有变量
        self._callbacks = callbacks
        self._helpers = callbacks.getHelpers()
        # 实现你想实现的代码
        callbacks.setExtensionName(self._extensionName)
        # 控制台标准输出、错误输出
        self._stdout = PrintWriter(callbacks.getStdout(), True)
        _stderr = PrintWriter(callbacks.getStderr(), True)
        # self._stdout.println("Hello Burp Mysql 1.0")
        install = time.strftime("%Y-%m-%d", time.localtime())
        self._stdout.println("+------------------------------+")
        self._stdout.println("|         Burp Mysql 1.0         |")
        self._stdout.println("|    Started @ " + install + "      |")
        self._stdout.println("+------------------------------+")
        # stderr.println("Hello erroutputs")
        self._log = ArrayList()  # java
        self._lock = Lock()
        self._mysql = None

        # 这里只定义一个界面Options界面，主要配置Mysql参数
        # https: // blog.csdn.net / xietansheng / article / details / 74366517
        self.mainTab = JTabbedPane()

        # options面板,垂直布局
        optionTop = JPanel()
        layout_top = BoxLayout(optionTop, BoxLayout.Y_AXIS)
        optionTop.setLayout(layout_top)
        optionTop.setBorder(BorderFactory.createTitledBorder("Mysql Configuration"))
        project_name = JTextField('project name: ')
        mysql_ip = JTextField('Mysql ip: ')
        mysql_username = JTextField('Mysql username: ')
        mysql_password = JTextField('Mysql password: ')
        self.connect_button = JButton("Connect")
        self.connect_button.addActionListener(connectMysql(self))
        optionTop.add(project_name)
        optionTop.add(mysql_ip)
        optionTop.add(mysql_username)
        optionTop.add(mysql_password)
        optionTop.add(self.connect_button)

        self.mainTab.add("Options", optionTop)

        # callbacks.registerProxyListener(self)
        # 美容UI
        self._callbacks.customizeUiComponent(self.mainTab)
        # 一定要加ITab,不然没有界面
        self._callbacks.addSuiteTab(self)
        # 注册httpListener
        self._callbacks.registerHttpListener(self)
        return

    '''
    void processHttpMessage(int toolFlag,
            boolean messageIsRequest,
            IHttpRequestResponse messageInfo);
    拦截HTTP请求
    '''

    def getBody(self, rawMessage, parsedMessage):
        # bodyOffset = requestInfo.getBodyOffset()
        # helpers中带bytes 转 string
        # bodyBytes = messageInfo.getRequest()[bodyOffset:]
        # bodyStrings = self._helpers.bytesToString(bodyBytes)
        # 将bytes转化成string
        return self._helpers.bytesToString(rawMessage[parsedMessage.getBodyOffset():])

    def processHttpMessage(self, toolFlag, messageIsRequest, messageInfo):

        # 如果按钮未开启，则退出
        if self.connect_button.getText() == 'Connect':
            return
        # 既有响应又有请求，很重要(不然会出现很大的bug)
        # 判断请求是否是PROXY中的
        if messageIsRequest:
            return
        else:
            # self._stdout.println("Enter print request")
            if toolFlag == 4:
                try:
                    # 使用requestinfo我们可以轻松的获得body和headers
                    requestInfo = self._helpers.analyzeRequest(messageInfo.getRequest())
                    # string
                    body = self.getBody(messageInfo.getRequest(), requestInfo)
                    if self._mysql != None:
                        self._stdout.println(body)
                        self._lock.acquire()
                        self._log = LogEntry(self,toolFlag, messageInfo, self._helpers, self._callbacks)
                        self._stdout.println(self._log._mime)
                        self._stdout.println(self._log._headers)
                        self._stdout.println(self._log._url)
                        self._stdout.println(self._log._method)
                        self._stdout.println(self._log.content_type)
                        self._stdout.println(self._log._path_params)
                        self._stdout.println(self._log._status)
                        insert_url_rec(self._mysql, self._log)
                        self._lock.release()
                    else:
                        return
                except Exception as e:
                    print("messageIsRequest error", e)
                    return

    # Give the new tab a name
    def getTabCaption(self):
        return self._labelName

    def getUiComponent(self):
        return self.mainTab


'''
public int getRowCount();
public int getColumnCount();
public Object getValueAt(int row, int column);
'''


# 创建log实体类，来记录每个请求（实际就是将请求给抽象成模型）
# __init__魔术方法，只是将传入的参数来初始化该实例
# __new__用来创建类并返回这个类的实例
class LogEntry:

    def __init__(self, extender,toolFlag, messageInfo, helpers, callbacks):
        self._extender=extender
        self._callbacks = callbacks
        self._helpers = helpers
        self._toolFlag = toolFlag
        self._requestResponse = self._callbacks.saveBuffersToTempFiles(messageInfo)
        if self._requestResponse.getResponse():
            responseInfo = self._helpers.analyzeResponse(self._requestResponse.getResponse())
            self._status = responseInfo.getStatusCode()
            self._mime = responseInfo.getStatedMimeType()
        requestInfo = self._helpers.analyzeRequest(messageInfo)
        # 分析request请求
        parsed_http_message = self._helpers.analyzeRequest(messageInfo.getRequest())
        self._url = requestInfo.getUrl()
        self._host = self._url.getHost()
        self._port = self._url.getPort()
        self._headers = parsed_http_message.getHeaders().toString()
        self._method = parsed_http_message.getMethod()
        self.content_type = requestInfo.getContentType()
        # 请求data提取
        self._body = self._extender.getBody(messageInfo.getRequest(), requestInfo)
        # 取出path路径
        url_parsed = urlparse(self._url.toString())
        path = url_parsed.path
        self._query = url_parsed.query
        self._path_params = path
        if self._query:
            self._path_params += '?' + self._query
        else:
            self._query =str(self._query)+" "
        #self._mime_type = parsed_http_message.getInferredMimeType()
        #self._port = self._url.getPort()
        self._httpService = self._requestResponse.getHttpService()
        if self._httpService.getProtocol() == "https":
            self._protocol = True
        else:
            self._protocol = False
        # 预制为空，只要做了fuzz,才会有
        self.fuzzpayload = "xxxx"
        self._project_id=2
